import tushare as ts
import pandas as pd
import numpy as np
import talib
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
from sklearn.model_selection import GridSearchCV
import matplotlib.font_manager as fm

font_path = fm.findfont(fm.FontProperties(family='SimHei'))
plt.rcParams['font.family'] = fm.FontProperties(fname=font_path).get_name()

# 设置 Tushare token
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 获取数据
df = pro.daily(ts_code='600036.SH', start_date='20200101', end_date='20241231')
df = df.sort_values('trade_date')
df = df.set_index('trade_date')

# 生成基础衍生变量
df['close-open'] = (df['close'] - df['open']) / df['open']
df['high-low'] = (df['high'] - df['low']) / df['low']
df['pre_close'] = df['close'].shift(1)
df['price_change'] = df['close'] - df['pre_close']
df['p_change'] = (df['close'] - df['pre_close']) / df['pre_close'] * 100

# 技术指标计算
df['MA5'] = df['close'].rolling(5).mean()
df['MA10'] = df['close'].rolling(10).mean()
df['RSI'] = talib.RSI(df['close'], timeperiod=14)
df['MOM'] = talib.MOM(df['close'], timeperiod=5)
df['EMA12'] = talib.EMA(df['close'], timeperiod=12)
df['EMA26'] = talib.EMA(df['close'], timeperiod=26)
macd, signal, hist = talib.MACD(df['close'], fastperiod=6, slowperiod=12, signalperiod=9)
df['MACD'] = macd
df['MACDsignal'] = signal
df['MACDhist'] = hist

# 删除缺失值
df.dropna(inplace=True)

# 查看列名，确认是否存在 volume 列
print(df.columns)

# 定义特征与目标变量
# 若不存在 volume 列，则从特征列表中移除
if 'volume' not in df.columns:
    features = ['close', 'close-open', 'high-low',
                'MA5', 'MA10', 'RSI', 'MOM', 'EMA12', 'EMA26',
                'MACD', 'MACDsignal', 'MACDhist', 'p_change']
else:
    features = ['close', 'volume', 'close-open', 'high-low',
                'MA5', 'MA10', 'RSI', 'MOM', 'EMA12', 'EMA26',
                'MACD', 'MACDsignal', 'MACDhist', 'p_change']

X = df[features]
y = np.where(df['price_change'].shift(-1) > 0, 1, -1)

# 划分训练集与测试集
split = int(len(X) * 0.9)
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]

# 建立随机森林模型
model = RandomForestClassifier(
    n_estimators=50,
    max_depth=5,
    min_samples_leaf=10,
    random_state=42
)
model.fit(X_train, y_train)

# 预测与准确率
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"模型准确率：{accuracy:.4f}")

# 特征重要性分析
importances = model.feature_importances_
feature_importance = pd.DataFrame({
    '特征': X.columns,
    '重要性': importances
}).sort_values(by='重要性', ascending=False)
print(feature_importance.head())

# 收益回测曲线
X_test['origin_return'] = (X_test['close'] / X_test['close'].shift(1)).cumprod()
X_test['prediction'] = y_pred
X_test['strategy_return'] = (X_test['prediction'] * X_test['p_change'] / 100 + 1).cumprod()

plt.figure(figsize=(12, 6))
plt.plot(X_test['origin_return'], label='原始收益率')
plt.plot(X_test['strategy_return'], label='模型策略收益率')
plt.title('收益回测曲线')
plt.xlabel('日期')
plt.ylabel('累积收益率')
plt.legend()
plt.grid(True)
plt.show()
